import json
from pathlib import Path
from typing import Any, Optional, Dict
import time

from lightning.pytorch.loggers.logger import Logger
from lightning.pytorch.utilities import rank_zero_only
from PIL import Image

# 日志文件夹的根路径
LOG_PATH = Path("outputs/local")


class LocalLogger(Logger):
    """
    一个本地日志记录器，它将日志以流式（streaming）方式写入JSON Lines (.jsonl) 文件。
    - 超参数保存到 hparams.json。
    - 指标在每一步都实时追加到 metrics.jsonl。
    - 图片元数据实时追加到 images.jsonl。
    - 每次运行都会创建一个带时间戳的版本目录。
    """
    def __init__(self, name: Optional[str] = None, version: Optional[str] = None):
        super().__init__()
        self._name = name or "default_experiment"
        self._version = version or str(int(time.time()))
        
        # 创建本次运行的专属日志目录
        self.log_dir.mkdir(parents=True, exist_ok=True)
        
        # 定义日志文件路径
        self.hparams_file = self.log_dir / "hparams.json"
        self.metrics_file = self.log_dir / "metrics.jsonl"
        self.images_file = self.log_dir / "images.jsonl"

    @property
    def name(self) -> str:
        return self._name

    @property
    def version(self) -> str:
        return self._version
    
    @property
    def log_dir(self) -> Path:
        """返回本次运行的日志目录，例如: outputs/local/default_experiment/1723260127"""
        return LOG_PATH / self.name / self.version

    @rank_zero_only
    def log_hyperparams(self, params: Dict[str, Any]) -> None:
        """记录超参数，一次性写入JSON文件"""
        with open(self.hparams_file, "w", encoding="utf-8") as f:
            json.dump(params, f, indent=4)

    @rank_zero_only
    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
        """记录指标，以追加模式写入 .jsonl 文件"""
        log_entry = {"step": step, **metrics}
        
        # 'a' 代表追加模式 (append)
        with open(self.metrics_file, "a", encoding="utf-8") as f:
            f.write(json.dumps(log_entry) + '\n')

    @rank_zero_only
    def log_image(
        self,
        key: str,
        images: list[Any],
        step: Optional[int] = None,
        **kwargs,
    ):
        """保存图片文件，并将元数据追加到 .jsonl 文件"""
        assert step is not None
        
        # 准备图片元数据列表
        image_metadata = []
        for index, image in enumerate(images):
            # 定义并创建图片保存路径
            image_path = self.log_dir / "images" / key
            image_path.mkdir(exist_ok=True, parents=True)
            
            # 保存图片文件
            filename = f"{index:0>2}_{step:0>6}.jpg"
            full_path = image_path / filename
            Image.fromarray(image).save(full_path)
            
            # 记录图片元数据
            image_metadata.append({
                "step": step,
                "key": key,
                "index": index,
                "path": str(full_path)
            })
            
        # 以追加模式写入图片日志
        with open(self.images_file, "a", encoding="utf-8") as f:
            for meta in image_metadata:
                f.write(json.dumps(meta) + '\n')

    @rank_zero_only
    def finalize(self, status: str) -> None:
        """在训练结束时调用，可以用来打印最终信息。"""
        print(f"Logger finalized. Logs are saved in: {self.log_dir}")